import time
from typing import List, Optional
from langchain_community.chat_models.baichuan import ChatBaichuan
from langchain_community.embeddings import BaichuanTextEmbeddings

from common.constants import *


class MyBaichuanTextEmbeddings(BaichuanTextEmbeddings):
    def embed_documents(self, texts: List[str]) -> Optional[List[List[float]]]:  # type: ignore[override]
        if texts and len(texts) > 15:
            result = []
            for i in range(len(texts) // 15 + 1):
                slice = texts[i * 15:i * 15 + 15]
                if not slice:
                    pass
                result += super()._embed(slice)
                print("loading {}/{}, {}".format(i, len(texts) // 15 + 1, int(time.time())))
                time.sleep(2)
            else:
                return result
        return super()._embed(texts)


llm = ChatBaichuan(baichuan_api_key=BAICHUAN_API_KEY)
embedding_model = MyBaichuanTextEmbeddings(baichuan_api_key=BAICHUAN_API_KEY)
